import json
import datasets
import pandas as pd
from transformers import AutoTokenizer
from PyCmpltrtok.common import get_dir_name_ext
import os


def json2limit(xpath, limit):
    assert(isinstance(limit, int) and limit > 0)

    print(f'Reading {xpath}')
    with open(xpath, 'r', encoding='utf8') as fi:
        xcont = fi.read()
        xlen = len(xcont)
        print(f'Length: {xlen:,d}')
        xdata = json.loads(xcont)

    cnt = -1
    xall_data = []
    for xobj in xdata:
        cnt += 1
        if cnt >= limit:
            break
        xall_data.append(xobj)

    xdir, xname, xext = get_dir_name_ext(xpath)
    xnew_path = os.path.join(xdir, xname + f'.limit-{limit}.' + xext)
    print(f'Saving to {xnew_path}')
    with open(xnew_path, 'w', encoding='utf8') as fo:
        json.dump(xall_data, fo)
    print('Done')


if '__main__' == __name__:

    def _main():
        xpath_arr = [
            r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeEE-V2\CMeEE-V2_dev.json',
            r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeEE-V2\CMeEE-V2_train.json',
        ]

        for xpath in xpath_arr:
            json2limit(xpath, 30)

    _main()
